package modbus

import (
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"time"
)

// ServerSession tcp server session
type ServerSession struct {
	conn         net.Conn
	readTimeout  time.Duration
	writeTimeout time.Duration
	*serverCommon
	logger
}

// handler net conn
func (sf *ServerSession) running(ctx context.Context) {
	var err error
	var bytesRead int

	sf.Debugf("client(%v) -> server(%v) connected", sf.conn.RemoteAddr(), sf.conn.LocalAddr())
	defer func() {
		// Ensure connection is always closed, even if panic occurs
		if closeErr := sf.conn.Close(); closeErr != nil {
			sf.Debugf("error closing connection: %v", closeErr)
		}
		sf.Debugf("client(%v) -> server(%v) disconnected, cause: %v", sf.conn.RemoteAddr(), sf.conn.LocalAddr(), err)
	}()

	raw := make([]byte, tcpAduMaxSize)
	for {
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return
		default:
		}

		adu := raw
		for rdCnt, length := 0, tcpHeaderMbapSize; rdCnt < length; {
			// Set read deadline for each read operation
			if deadlineErr := sf.conn.SetReadDeadline(time.Now().Add(sf.readTimeout)); deadlineErr != nil {
				err = deadlineErr
				return
			}

			bytesRead, err = io.ReadFull(sf.conn, adu[rdCnt:length])
			if err != nil {
				// Improved error handling
				if sf.isConnectionClosed(err) {
					return
				}

				if netErr, ok := err.(net.Error); ok {
					if netErr.Timeout() {
						err = fmt.Errorf("read timeout: %v", err)
						return
					}
					if !netErr.Temporary() {
						return
					}
				}

				if bytesRead == 0 && err == io.EOF {
					err = fmt.Errorf("remote client closed: %v", err)
					return
				}
				// For other errors, continue with next iteration
			}

			rdCnt += bytesRead
			if rdCnt >= length {
				// Check head ProtocolIdentifier
				if binary.BigEndian.Uint16(adu[2:]) != tcpProtocolIdentifier {
					rdCnt, length = 0, tcpHeaderMbapSize
					continue
				}
				length = int(binary.BigEndian.Uint16(adu[4:])) + tcpHeaderMbapSize - 1
				// Validate message length to prevent buffer overflow
				if length > tcpAduMaxSize {
					err = fmt.Errorf("message too large: %d bytes", length)
					return
				}
				if rdCnt == length {
					if frameErr := sf.frameHandler(adu[:length]); frameErr != nil {
						err = frameErr
						return
					}
					// Reset for next message
					rdCnt, length = 0, tcpHeaderMbapSize
				}
			}
		}
	}
}

// isConnectionClosed checks if the error indicates a closed connection
func (sf *ServerSession) isConnectionClosed(err error) bool {
	if err == io.EOF || err == io.ErrClosedPipe {
		return true
	}

	if strings.Contains(err.Error(), "use of closed network connection") {
		return true
	}

	if strings.Contains(err.Error(), "connection reset by peer") {
		return true
	}

	return false
}

// modbus 包处理
func (sf *ServerSession) frameHandler(requestAdu []byte) error {
	defer func() {
		if err := recover(); err != nil {
			sf.Errorf("painc happen,%v", err)
		}
	}()

	sf.Debugf("RX Raw[% x]", requestAdu)
	// got head from request adu
	tcpHeader := protocolTCPHeader{
		binary.BigEndian.Uint16(requestAdu[0:]),
		binary.BigEndian.Uint16(requestAdu[2:]),
		binary.BigEndian.Uint16(requestAdu[4:]),
		requestAdu[6],
	}
	funcCode := requestAdu[7]
	pduData := requestAdu[8:]

	node, err := sf.GetNode(tcpHeader.slaveID)
	if err != nil { // slave id not exit, ignore it
		return nil
	}
	var rspPduData []byte
	if handle, ok := sf.function[funcCode]; ok {
		rspPduData, err = handle(node, pduData)
	} else {
		err = &ExceptionError{ExceptionCodeIllegalFunction}
	}
	if err != nil {
		funcCode |= 0x80
		rspPduData = []byte{err.(*ExceptionError).ExceptionCode}
	}

	// prepare responseAdu data,fill it
	responseAdu := requestAdu[:tcpHeaderMbapSize]
	binary.BigEndian.PutUint16(responseAdu[0:], tcpHeader.transactionID)
	binary.BigEndian.PutUint16(responseAdu[2:], tcpHeader.protocolID)
	binary.BigEndian.PutUint16(responseAdu[4:], uint16(2+len(rspPduData)))
	responseAdu[6] = tcpHeader.slaveID
	responseAdu = append(responseAdu, funcCode)
	responseAdu = append(responseAdu, rspPduData...)

	sf.Debugf("TX Raw[% x]", responseAdu)
	// write response
	return sf.writeResponse(responseAdu)
}

// writeResponse writes the response with proper error handling
func (sf *ServerSession) writeResponse(responseAdu []byte) error {
	for wrCnt := 0; len(responseAdu) > wrCnt; {
		// Set write deadline for each write operation
		if err := sf.conn.SetWriteDeadline(time.Now().Add(sf.writeTimeout)); err != nil {
			return fmt.Errorf("set write deadline: %v", err)
		}

		byteCount, err := sf.conn.Write(responseAdu[wrCnt:])
		if err != nil {
			// Check if connection is closed
			if sf.isConnectionClosed(err) {
				return err
			}

			if netErr, ok := err.(net.Error); ok {
				if netErr.Timeout() {
					return fmt.Errorf("write timeout: %v", err)
				}
				if netErr.Temporary() {
					// For temporary errors, we can retry
					continue
				}
			}
			return err
		}
		wrCnt += byteCount
	}
	return nil
}
